{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f222b04a-83fa-431b-a6b3-77afca19c674",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "import jax_md\n",
    "import numpy as np\n",
    "import pathlib\n",
    "\n",
    "from ase.io import read\n",
    "from jax_md import units\n",
    "from typing import Dict\n",
    "\n",
    "from so3lr import to_jax_md\n",
    "from so3lr import So3lrPotential\n",
    "\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "c0886332-077a-4312-88ed-f9c7d115772f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# If you want to perform simulations in float64 you have to call this before any JAX compuation\n",
    "jax.config.update('jax_enable_x64', True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f9b081c-1f15-4c85-87c0-ec0804e09e6e",
   "metadata": {},
   "source": [
    "# Nose-Hoover NPT\n",
    "\n",
    "The code below is the code already introduces in the `npt_jaxmd.ipynb` example. For details check it out. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "2302bbec-2d97-4aa7-b985-ba5875776896",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simulation parameters\n",
    "\n",
    "timestep = 0.0005  # Time step in ps\n",
    "npt_cycles = 25  # Number of Cycles in the NVT.\n",
    "npt_steps = 5  # Number of NVT steps per cylce. The total number of MD steps equals nvt_cylces * nvt_steps\n",
    "\n",
    "T_init = 300  # Initial temperature.\n",
    "pressure = 1.  # Target pressure. \n",
    "\n",
    "chain = 3  # Number of chains in the Nose-Hoover chain.\n",
    "chain_steps = 2  # Number of steps per chain.\n",
    "sy_steps = 3\n",
    "\n",
    "thermo = 100  # Thermostat value in the Nose-Hoover chain. \n",
    "baro = 1000  # Barostat value in the Nose-Hoover chain. \n",
    "\n",
    "# Default nose hoover chain parameters.\n",
    "def default_nhc_kwargs(\n",
    "    tau: jnp.float32, \n",
    "    overrides: Dict\n",
    ") -> Dict:\n",
    "    \n",
    "    default_kwargs = {\n",
    "        'chain_length': 3, \n",
    "        'chain_steps': 2, \n",
    "        'sy_steps': 3,\n",
    "        'tau': tau\n",
    "    }\n",
    "    \n",
    "    if overrides is None:\n",
    "        return default_kwargs\n",
    "  \n",
    "    return {\n",
    "      k: overrides.get(k, default_kwargs[k]) for k in default_kwargs\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b60d4845-aad3-4222-a6b3-66e87c0a5463",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "WARNING:absl:CheckpointMetadata file does not exist: /Users/thorbenfrank/Documents/git/so3lr/so3lr/params/checkpoints/ckpt_5150000/_CHECKPOINT_METADATA\n",
      "WARNING:absl:`StandardCheckpointHandler` expects a target tree to be provided for restore. Not doing so is generally UNSAFE unless you know the present topology to be the same one as the checkpoint was saved under.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "jax-md partition: cell_size=0.1814516129032258, cl.id_buffer.size=2875, N=1536, cl.did_buffer_overflow=False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/thorbenfrank/Documents/venvs/so3lr/lib/python3.12/site-packages/jax/_src/ops/scatter.py:92: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=int64 to dtype=int32 with jax_numpy_dtype_promotion='standard'. In future JAX releases this will result in an error.\n",
      "  warnings.warn(\n",
      "/Users/thorbenfrank/Documents/venvs/so3lr/lib/python3.12/site-packages/jax/_src/ops/scatter.py:92: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=int64 to dtype=int32 with jax_numpy_dtype_promotion='standard'. In future JAX releases this will result in an error.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step\tE\tFmax\n",
      "--------------------------\n",
      "jax-md partition: cell_size=0.1814516129032258, cl.id_buffer.size=2875, N=1536, cl.did_buffer_overflow=Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=3/0)>\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/thorbenfrank/Documents/venvs/so3lr/lib/python3.12/site-packages/jax/_src/ops/scatter.py:92: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=int64 to dtype=int32 with jax_numpy_dtype_promotion='standard'. In future JAX releases this will result in an error.\n",
      "  warnings.warn(\n",
      "/Users/thorbenfrank/Documents/venvs/so3lr/lib/python3.12/site-packages/jax/_src/ops/scatter.py:92: FutureWarning: scatter inputs have incompatible types: cannot safely cast value from dtype=int64 to dtype=int32 with jax_numpy_dtype_promotion='standard'. In future JAX releases this will result in an error.\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\t-2781.35\t1.23\n",
      "1\t-2819.48\t1.91\n"
     ]
    }
   ],
   "source": [
    "# Read some molecular structure. Here we read the water box from the test directory.\n",
    "path_to_xyz = '../tests/test_data/water_64.xyz'\n",
    "atoms = read(path_to_xyz, index=-1)\n",
    "atoms.wrap()\n",
    "\n",
    "# Repeat in each direction.\n",
    "atoms = atoms * [2, 2, 2]\n",
    "\n",
    "species = jnp.array(atoms.get_atomic_numbers())\n",
    "masses = jnp.array(atoms.get_masses())\n",
    "num_atoms = len(species)\n",
    "\n",
    "# it is important that scalar is array of dim = 1. Otherwise weird behavior of jax_md.simulate.npt_box(state) which returns\n",
    "# a 3x3 array when passed a jnp.array of ndim = 0. This leads to TracerConversionError. Maybe open issue on jax_md?\n",
    "\n",
    "box = jnp.array(\n",
    "    [\n",
    "        np.array(atoms.get_cell())[0, 0]\n",
    "    ]\n",
    ")  \n",
    "\n",
    "fractional_coordinates = True\n",
    "displacement, shift = jax_md.space.periodic_general(box=box, fractional_coordinates=fractional_coordinates)\n",
    "\n",
    "positions_init = jnp.array(atoms.get_positions())\n",
    "if fractional_coordinates:\n",
    "    positions_init = positions_init / box\n",
    "\n",
    "neighbor_fn, neighbor_fn_lr, energy_fn = to_jax_md(\n",
    "    potential=So3lrPotential(\n",
    "        dtype=jnp.float64  # or float32 for single precision\n",
    "    ),\n",
    "    displacement_or_metric=displacement,\n",
    "    box_size=box,\n",
    "    species=species,\n",
    "    capacity_multiplier=1.25,\n",
    "    buffer_size_multiplier_sr=1.25,\n",
    "    buffer_size_multiplier_lr=1.25,\n",
    "    minimum_cell_size_multiplier_sr=1.0,\n",
    "    disable_cell_list=False,\n",
    "    fractional_coordinates=fractional_coordinates\n",
    ")\n",
    "\n",
    "# Energy function.\n",
    "energy_fn = jax.jit(energy_fn)\n",
    "force_fn = jax.jit(jax_md.quantity.force(energy_fn))\n",
    "\n",
    "# Initialize the short and long-range neighbor lists.\n",
    "nbrs = neighbor_fn.allocate(\n",
    "    positions_init,\n",
    "    box=box\n",
    ")\n",
    "nbrs_lr = neighbor_fn_lr.allocate(\n",
    "    positions_init,\n",
    "    box=box\n",
    ")\n",
    "\n",
    "# For repeated execution of cell, delete the fire_state first, otherwise jit in jupyter environment \n",
    "# can meddle with things.\n",
    "try:\n",
    "    del fire_state\n",
    "except NameError:\n",
    "    pass\n",
    "\n",
    "min_cycles = 5\n",
    "min_steps = 10\n",
    "\n",
    "fire_init, fire_apply = jax_md.minimize.fire_descent(\n",
    "    energy_fn, \n",
    "    shift, \n",
    "    dt_start = 0.05, \n",
    "    dt_max = 0.1, \n",
    "    n_min = 2\n",
    ")\n",
    "\n",
    "fire_apply = jax.jit(fire_apply)\n",
    "fire_state = fire_init(\n",
    "    positions_init, \n",
    "    box=box, \n",
    "    neighbor=nbrs.idx,\n",
    "    neighbor_lr=nbrs_lr.idx\n",
    ")\n",
    "\n",
    "@jax.jit\n",
    "def step_fire_fn(i, fire_state):\n",
    "    \n",
    "    fire_state, nbrs, nbrs_lr = fire_state\n",
    "    \n",
    "    fire_state = fire_apply(\n",
    "        fire_state, \n",
    "        neighbor=nbrs.idx, \n",
    "        neighbor_lr=nbrs_lr.idx, \n",
    "        box=box\n",
    "    )\n",
    "    \n",
    "    nbrs = nbrs.update(\n",
    "        fire_state.position,\n",
    "        neighbor=nbrs.idx,\n",
    "        box=box\n",
    "    )\n",
    "    \n",
    "    nbrs_lr = nbrs_lr.update(\n",
    "        fire_state.position,\n",
    "        neighbor_lr=nbrs_lr.idx,\n",
    "        box=box\n",
    "    )\n",
    "    \n",
    "    return fire_state, nbrs, nbrs_lr\n",
    "\n",
    "print('Step\\tE\\tFmax')\n",
    "print('--------------------------')\n",
    "for i in range(min_cycles):\n",
    "    fire_state, nbrs, nbrs_lr = jax.lax.fori_loop(\n",
    "        0, \n",
    "        min_steps, \n",
    "        step_fire_fn, \n",
    "        (fire_state, nbrs, nbrs_lr)\n",
    "    )\n",
    "    \n",
    "    E = energy_fn(\n",
    "        fire_state.position, \n",
    "        neighbor=nbrs.idx, \n",
    "        neighbor_lr=nbrs_lr.idx,\n",
    "        box=box\n",
    "    )\n",
    "\n",
    "    F = force_fn(\n",
    "        fire_state.position, \n",
    "        neighbor=nbrs.idx, \n",
    "        neighbor_lr=nbrs_lr.idx,\n",
    "        box=box\n",
    "    )\n",
    "    \n",
    "    print('{}\\t{:.2f}\\t{:.2f}'.format(i, E, np.abs(F).max()))\n",
    "\n",
    "# Dictionary with the NHC settings.\n",
    "new_nhc_kwargs = {\n",
    "    'chain_length': chain, \n",
    "    'chain_steps': chain_steps, \n",
    "    'sy_steps': sy_steps\n",
    "}\n",
    "\n",
    "# Convert to metal unit system.\n",
    "unit = units.metal_unit_system()\n",
    "\n",
    "timestep = timestep * unit['time']\n",
    "T_init = T_init * unit['temperature']\n",
    "\n",
    "rng_key = jax.random.PRNGKey(0)\n",
    "\n",
    "# Choose Nose-Hoover thermostat.\n",
    "init_fn, apply_fn = jax_md.simulate.npt_nose_hoover(\n",
    "    energy_fn, \n",
    "    shift, \n",
    "    dt=timestep,\n",
    "    pressure=pressure, \n",
    "    kT=T_init,\n",
    "    barostat_kwargs=default_nhc_kwargs(baro * timestep, new_nhc_kwargs),\n",
    "    thermostat_kwargs=default_nhc_kwargs(thermo * timestep, new_nhc_kwargs)\n",
    ")\n",
    "\n",
    "apply_fn = jax.jit(apply_fn)\n",
    "init_fn = jax.jit(init_fn)\n",
    "\n",
    "# Initialize state using position and neigbhors structure relaxation.\n",
    "state = init_fn(\n",
    "    rng_key, \n",
    "    fire_state.position, \n",
    "    box=box, \n",
    "    neighbor=nbrs.idx, \n",
    "    neighbor_lr=nbrs_lr.idx,\n",
    "    kT=T_init,\n",
    "    mass=masses\n",
    ")\n",
    "\n",
    "@jax.jit\n",
    "def step_npt_fn(i, state):\n",
    "    state, nbrs, nbrs_lr, box = state\n",
    "    \n",
    "    state = apply_fn(\n",
    "        state,\n",
    "        neighbor=nbrs.idx, \n",
    "        neighbor_lr=nbrs_lr.idx, \n",
    "        kT=T_init,\n",
    "        pressure=pressure\n",
    "    )\n",
    "    \n",
    "    box = jax_md.simulate.npt_box(state)\n",
    "    \n",
    "    nbrs = nbrs.update(\n",
    "        state.position, \n",
    "        neighbor=nbrs.idx, \n",
    "        box=box\n",
    "    )\n",
    "    \n",
    "    nbrs_lr = nbrs_lr.update(\n",
    "        state.position,\n",
    "        neighbor_lr=nbrs_lr.idx,\n",
    "        box=box\n",
    "    )\n",
    "    \n",
    "    return state, nbrs, nbrs_lr, box"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd85e38b-52ce-4ddd-af0f-cc8480c3f35a",
   "metadata": {},
   "source": [
    "# Production run\n",
    "Production runs require efficient saving of the frames and additional properties during the simulations. Operations like `ase.io.write` can become very slow for larger structures with thousands of atoms and pose a bottleneck for efficient simulations in requiring orders of magnitude more time than the MD steps itself. This is, why for production runs we use `.hdf5` to allow for efficient writing to disk. \n",
    "\n",
    "It is build around `DataSetEntry` and `HDF5Store`, provided by the `mlff` package. Below you find an example function called `init_hdf5_store` which initializes and `.hdf5` data base under `save_to`. Instead of saving frames subsequently it is intended to first collect multiple values for the property of interest (e.g. multiple frames) and flush them in a single call to the data base. The number of values that is flushed at once is the `batch_size`. As the data base needs to know the shapes a priori, each property that is intended to be saved needs to be initialized with the correct shapes. E.g. for the positions its `(batch_size, num_atoms, 3)`. During simulation arrays of the specified size will be stacked along the 0-th axis, such that after e.g. 10 flushing operations `positions` will be of size `(10, batch_size, num_atoms, 3)`. In addition to positions we will save the `velocities` and the `box` which can be represented in terms of a fixed number of entries. For details on the box representation in `jax_md` see the `npt_jaxmd.ipynb` example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2ff24970-d203-4761-8650-c1a78e9e2a25",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlff.mdx.hdfdict import DataSetEntry, HDF5Store\n",
    "\n",
    "\n",
    "def init_hdf5_store(\n",
    "    save_to, \n",
    "    batch_size, \n",
    "    num_atoms,\n",
    "    num_box_entries,\n",
    "    exist_ok=False\n",
    "):\n",
    "    \n",
    "    parent_dir = pathlib.Path(save_to).expanduser().resolve().parent\n",
    "    parent_dir.mkdir(exist_ok=True)\n",
    "    \n",
    "    _save_to = pathlib.Path(save_to)\n",
    "    if _save_to.exists():\n",
    "        if exist_ok is False:\n",
    "            raise RuntimeError(\n",
    "                f'File exists save_to={_save_to}. '\n",
    "                f'Set exists_ok=True to override file.'\n",
    "            )\n",
    "    \n",
    "    dataset = {\n",
    "        'positions': DataSetEntry(\n",
    "            chunk_length=1, \n",
    "            shape=(batch_size, num_atoms, 3), \n",
    "            dtype=np.float32\n",
    "        ),\n",
    "        'velocities': DataSetEntry(\n",
    "            chunk_length=1, \n",
    "            shape=(batch_size, num_atoms, 3), \n",
    "            dtype=np.float32\n",
    "        ),\n",
    "        'box': DataSetEntry(\n",
    "            chunk_length=1, \n",
    "            shape=(batch_size, num_box_entries), \n",
    "            dtype=np.float32\n",
    "        )\n",
    "    }\n",
    "\n",
    "    return HDF5Store(_save_to, datasets=dataset, mode='w')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "f422d2bf-83fd-4ae7-a35b-7d6f3365db33",
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the hdf5_store object.\n",
    "\n",
    "# How many npt_cicle outputs are collected before flush is called. \n",
    "# 100 is a good default for longer runs.\n",
    "hdf5_buffer_size = 3\n",
    "\n",
    "hdf5_store = init_hdf5_store(\n",
    "    save_to='hdf5/trajectory.hdf5',\n",
    "    batch_size=hdf5_buffer_size,  # number of frames that is flushed in one call\n",
    "    num_atoms=num_atoms,  # number of atoms\n",
    "    num_box_entries=1,  # cubic box is represented by the 3 values in jax-md\n",
    "    exist_ok=False  # if save_to is allowed to exist. If true, the .hdf5 file will be overwritten\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a098e470-ed61-4daf-9720-ed7b8d6b179d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Step\tKE\tPE\tTot. Energy\tTemp.\ttime/steps\n",
      "------------------------------------------------------------------------------------------------------------\n",
      "jax-md partition: cell_size=0.18145161867141724, cl.id_buffer.size=2875, N=1536, cl.did_buffer_overflow=Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=3/0)>\n",
      "5\t42.14\t-5114.97\t-5072.826\t212.2\t10.79\n",
      "10\t45.97\t-5119.05\t-5073.082\t231.5\t2.57\n",
      "15\t36.36\t-5109.16\t-5072.802\t183.1\t2.64\n",
      "20\t41.95\t-5114.77\t-5072.830\t211.3\t2.55\n",
      "25\t35.00\t-5107.59\t-5072.594\t176.3\t2.63\n",
      "30\t31.86\t-5104.22\t-5072.366\t160.5\t2.57\n",
      "35\t39.47\t-5111.77\t-5072.297\t198.8\t2.56\n",
      "40\t39.45\t-5111.31\t-5071.858\t198.7\t2.56\n",
      "45\t41.33\t-5113.09\t-5071.765\t208.2\t2.61\n",
      "50\t35.43\t-5106.94\t-5071.507\t178.5\t2.56\n",
      "55\t40.58\t-5111.99\t-5071.410\t204.4\t2.60\n",
      "60\t43.20\t-5114.41\t-5071.207\t217.6\t2.58\n",
      "65\t39.56\t-5110.55\t-5070.987\t199.3\t2.56\n",
      "70\t41.13\t-5112.06\t-5070.927\t207.2\t2.60\n",
      "75\t41.99\t-5112.61\t-5070.623\t211.5\t2.57\n",
      "80\t47.53\t-5118.00\t-5070.471\t239.4\t2.57\n",
      "85\t43.08\t-5113.26\t-5070.180\t217.0\t2.57\n",
      "90\t42.56\t-5112.53\t-5069.967\t214.4\t2.58\n",
      "95\t46.05\t-5115.76\t-5069.714\t231.9\t2.57\n",
      "100\t46.80\t-5116.16\t-5069.365\t235.7\t2.57\n",
      "105\t47.28\t-5116.42\t-5069.140\t238.2\t2.56\n",
      "110\t44.44\t-5113.20\t-5068.761\t223.8\t2.63\n",
      "115\t49.71\t-5118.23\t-5068.516\t250.4\t2.58\n",
      "120\t50.35\t-5118.48\t-5068.135\t253.6\t2.58\n",
      "125\t48.70\t-5116.48\t-5067.787\t245.3\t2.56\n",
      "Total_time:  370.22679591178894\n"
     ]
    }
   ],
   "source": [
    "# Track total time and step times averaged over cycle.\n",
    "total_time = time.time()\n",
    "\n",
    "velocities, positions, boxes = [], [], []\n",
    "print('Step\\tKE\\tPE\\tTot. Energy\\tTemp.\\ttime/steps')\n",
    "print('------------------------------------------------------------------------------------------------------------')\n",
    "for i in range(npt_cycles):\n",
    "\n",
    "    old_time = time.time()\n",
    "    \n",
    "    # Do `npt_steps` NPT steps.\n",
    "    new_state, nbrs, nbrs_lr, new_box = jax.block_until_ready(\n",
    "        jax.lax.fori_loop(\n",
    "            0, \n",
    "            npt_steps, \n",
    "            step_npt_fn, \n",
    "            (state, nbrs, nbrs_lr, box)\n",
    "        )\n",
    "    )\n",
    "    \n",
    "    new_time = time.time()\n",
    "    \n",
    "    # Check for overflor of both sr and lr neighbors.\n",
    "    if nbrs.did_buffer_overflow:\n",
    "        print('Neighbor list overflowed, reallocating.')\n",
    "        nbrs = neighbor_fn.allocate(state.position, box = box)\n",
    "        if nbrs_lr.did_buffer_overflow:\n",
    "            print('Long-range neighbor list also overflowed, reallocating.')\n",
    "            nbrs_lr = neighbor_fn_lr.allocate(state.position, box = box)\n",
    "    elif nbrs_lr.did_buffer_overflow:\n",
    "        print('Long-range neighbor list overflowed, reallocating.')\n",
    "        nbrs_lr = neighbor_fn_lr.allocate(state.position, box = box)\n",
    "    else:\n",
    "        state = new_state\n",
    "        box = new_box\n",
    "\n",
    "        # Only save if there are no overflows in the neighbor lists.\n",
    "        \n",
    "        ######## HDF5 saving part ########\n",
    "        \n",
    "        # Append velocities, calculated from the momenta.\n",
    "        velocities.append(\n",
    "            jnp.divide(\n",
    "                state.momentum,\n",
    "                state.mass\n",
    "            )\n",
    "        )\n",
    "\n",
    "        # Append positions of the last frame of `npt_steps` MD steps.\n",
    "        positions.append(\n",
    "            jax_md.space.transform(box=box, R=state.position)  # this is only correct for cubic cells with equally long lattice vectors\n",
    "        )\n",
    "        \n",
    "        # Append current box.\n",
    "        boxes.append(\n",
    "            box\n",
    "        )\n",
    "    \n",
    "    # If buffer is full and positions is not empty, append velocities, positions and boxes for in \n",
    "    # `hdf5_buffer_size` frames to the trajectory.hdf5\n",
    "    if (len(positions) % hdf5_buffer_size == 0) and (len(positions) > 0):   \n",
    "        \n",
    "        # Create dictionary from stacked velocities and positions.\n",
    "        step_data = dict(\n",
    "            velocities=jnp.stack(velocities, axis=0),  # (hdf5_buffer_size, num_atoms, 3)\n",
    "            positions=jnp.stack(positions, axis=0),  # (hdf5_buffer_size, num_atoms, 3)\n",
    "            box=jnp.stack(boxes, axis=0)            # (hdf5_buffer_size, num_box_entries)\n",
    "        )\n",
    "        \n",
    "        # Make the jax.numpy arrays to numpy arrays.\n",
    "        step_data = jax.tree.map(lambda u: np.asarray(u), step_data)\n",
    "\n",
    "        # Append the step data to the hdf5 store object.\n",
    "        hdf5_store.append(\n",
    "            step_data\n",
    "        )\n",
    "        \n",
    "        # Create fresh lists.\n",
    "        velocities, positions, boxes = [], [], []\n",
    "    \n",
    "        ######## End of HDF5 saving part ########\n",
    "\n",
    "    # Calculate some quantities for printing\n",
    "    KE = jax_md.quantity.kinetic_energy(\n",
    "        momentum=state.momentum,\n",
    "        mass=state.mass\n",
    "    )\n",
    "    \n",
    "    PE = energy_fn(\n",
    "        state.position,\n",
    "        neighbor=nbrs.idx,\n",
    "        neighbor_lr=nbrs_lr.idx, \n",
    "        box=box\n",
    "    )\n",
    "    \n",
    "    T = jax_md.quantity.temperature(\n",
    "        momentum=state.momentum,\n",
    "        mass = state.mass\n",
    "    ) / unit['temperature']\n",
    "\n",
    "    print(\n",
    "        f'{(i+1)*npt_steps}\\t{KE:.2f}\\t{PE:.2f}\\t{KE+PE:.3f}\\t{T:.1f}\\t{(new_time - old_time) / npt_steps:.2f}'\n",
    "    )\n",
    "\n",
    "print('Total_time: ', time.time() - total_time)\n",
    "\n",
    "# Clear all caches\n",
    "jax.clear_caches()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a87b695f-b712-4b97-a185-127586d26c96",
   "metadata": {},
   "source": [
    "# Postprocess the Trajectory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "169f1144-4f94-489d-a77b-35feb8a9fb21",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(9, 3, 1536, 3)\n",
      "(27, 1536, 3)\n"
     ]
    }
   ],
   "source": [
    "import h5py\n",
    "\n",
    "file = h5py.File('hdf5/trajectory.hdf5', mode='r')\n",
    "\n",
    "# Load velocities and reshape.\n",
    "velocities = np.array(file['velocities'])\n",
    "# velocities have shape (num_flushes, hdf5_buffer_size, num_atoms, 3)\n",
    "print(velocities.shape)\n",
    "\n",
    "# reshape velocities such that one has only one leading axis of consecutive velocities\n",
    "# separated by delta_t = npt_steps * time_step [ps].\n",
    "velocities = velocities.reshape(-1, *file['velocities'].shape[-2:])\n",
    "\n",
    "print(velocities.shape)  # (num_flushes*hdf5_buffer_size, num_atoms, 3)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "so3lr",
   "language": "python",
   "name": "so3lr"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
